from typing import Dict, Any, List, Optional, TypedDict
from langgraph.graph import StateGraph, END
from agents import (
    ProductManagerAgent, ArchitectAgent, DeveloperAgent, 
    TesterAgent, ReviewerAgent, DocumenterAgent, 
    OptimizerAgent, SecurityCheckerAgent
)
from utils import logger, create_progress_bar, display_results_table
import time
import os
import json
import re

class WorkflowState(TypedDict):
    """工作流状态类型定义"""
    user_input: str
    original_requirement: str
    data: Dict[str, Any]
    results: Dict[str, Any]
    errors: List[Dict[str, Any]]
    execution_log: List[Dict[str, Any]]
    current_step: str
    completed_steps: List[str]

class CodeBuddyState:
    """CodeBuddy 工作流状态"""
    
    def __init__(self, initial_data: Dict[str, Any] = None):
        self.data = initial_data or {}
        self.results = {}
        self.errors = []
        self.execution_log = []
        self.current_step = ""
        self.start_time = time.time()
        # 添加完成状态跟踪 - 使用列表保持一致性
        self.completed_steps = []
    
    def update(self, key: str, value: Any):
        """更新状态数据"""
        self.data[key] = value
    
    def add_result(self, step: str, result: Dict[str, Any]):
        """添加步骤结果"""
        self.results[step] = result
        self.data.update(result)  # 将结果添加到数据流中
        if step not in self.completed_steps:  # 避免重复添加
            self.completed_steps.append(step)
        
        # 记录执行日志
        self.execution_log.append({
            'step': step,
            'timestamp': time.time(),
            'success': result.get('success', True),
            'summary': result.get('summary', '')
        })
    
    def add_error(self, step: str, error: str):
        """添加错误信息"""
        self.errors.append({
            'step': step,
            'error': error,
            'timestamp': time.time()
        })
    
    def is_step_completed(self, step: str) -> bool:
        """检查步骤是否完成"""
        return step in self.completed_steps
    
    def can_proceed_to_optimization(self) -> bool:
        """检查是否可以进入优化阶段"""
        return self.is_step_completed('review') or self.is_step_completed('security_check')
    
    def get_summary(self) -> Dict[str, Any]:
        """获取执行摘要"""
        successful_steps = sum(1 for log in self.execution_log if log['success'])
        total_steps = len(self.execution_log)
        execution_time = time.time() - self.start_time
        
        return {
            'total_steps': total_steps,
            'successful_steps': successful_steps,
            'failed_steps': total_steps - successful_steps,
            'execution_time': round(execution_time, 2),
            'success_rate': (successful_steps / total_steps * 100) if total_steps > 0 else 0,
            'errors': len(self.errors)
        }

class CodeBuddyWorkflow:
    """CodeBuddy 多智能体工作流"""
    
    def __init__(self):
        self.agents = self._initialize_agents()
        self.graph = self._build_graph()
    
    def _initialize_agents(self) -> Dict[str, Any]:
        """初始化所有AI代理"""
        agents = {
            'product_manager': ProductManagerAgent(),
            'architect': ArchitectAgent(),
            'developer': DeveloperAgent(),
            'tester': TesterAgent(),
            'reviewer': ReviewerAgent(),
            'documenter': DocumenterAgent(),
            'optimizer': OptimizerAgent(),
            'security_checker': SecurityCheckerAgent()
        }
        # 添加全局统计计数器
        self._execution_stats = {
            'total_calls': 0,
            'success_calls': 0,
            'total_tokens': 0
        }
        logger.info(f"✅ 已初始化 {len(agents)} 个AI代理")
        return agents

    def _update_execution_stats(self, result: Dict[str, Any]):
        """更新全局执行统计"""
        self._execution_stats['total_calls'] += 1
        if result.get('success', False):
            self._execution_stats['success_calls'] += 1
        self._execution_stats['total_tokens'] += result.get('tokens_used', 0)

    def _build_graph(self) -> StateGraph:
        """构建LangGraph状态图"""
        
        # 定义节点函数
        def product_analysis_node(state: Dict[str, Any]) -> Dict[str, Any]:
            """产品分析节点"""
            state['current_step'] = "需求分析"
            logger.info("🎯 开始需求分析...")
            
            try:
                result = self.agents['product_manager'].process(state['data'])
                result['success'] = True
                state['results']['product_analysis'] = result
                state['data'].update(result)
                state['completed_steps'].append('product_analysis')
                logger.info("✅ 需求分析完成")
            except Exception as e:
                state['errors'].append({
                    'step': 'product_analysis',
                    'error': str(e),
                    'timestamp': time.time()
                })
                state['results']['product_analysis'] = {'success': False, 'error': str(e)}
                logger.error(f"❌ 需求分析失败: {e}")
            
            return state
        
        def architecture_design_node(state: Dict[str, Any]) -> Dict[str, Any]:
            """架构设计节点"""
            state['current_step'] = "架构设计"
            logger.info("🏗️ 开始架构设计...")
            
            try:
                result = self.agents['architect'].process(state['data'])
                result['success'] = True
                state['results']['architecture_design'] = result
                state['data'].update(result)
                state['completed_steps'].append('architecture_design')
                logger.info("✅ 架构设计完成")
            except Exception as e:
                state['errors'].append({
                    'step': 'architecture_design',
                    'error': str(e),
                    'timestamp': time.time()
                })
                state['results']['architecture_design'] = {'success': False, 'error': str(e)}
                logger.error(f"❌ 架构设计失败: {e}")
            
            return state
        
        def development_node(state: Dict[str, Any]) -> Dict[str, Any]:
            """开发节点"""
            state['current_step'] = "代码开发"
            logger.info("💻 开始代码开发...")
            
            try:
                result = self.agents['developer'].process(state['data'])
                result['success'] = True
                state['results']['development'] = result
                state['data'].update(result)
                state['completed_steps'].append('development')
                logger.info("✅ 代码开发完成")
            except Exception as e:
                state['errors'].append({
                    'step': 'development',
                    'error': str(e),
                    'timestamp': time.time()
                })
                state['results']['development'] = {'success': False, 'error': str(e)}
                logger.error(f"❌ 代码开发失败: {e}")
            
            return state
        
        def testing_node(state: Dict[str, Any]) -> Dict[str, Any]:
            """测试节点"""
            state['current_step'] = "代码测试"
            logger.info("🧪 开始代码测试...")
            
            try:
                result = self.agents['tester'].process(state['data'])
                result['success'] = True
                state['results']['testing'] = result
                state['data'].update(result)
                state['completed_steps'].append('testing')
                logger.info("✅ 代码测试完成")
            except Exception as e:
                state['errors'].append({
                    'step': 'testing',
                    'error': str(e),
                    'timestamp': time.time()
                })
                state['results']['testing'] = {'success': False, 'error': str(e)}
                logger.error(f"❌ 代码测试失败: {e}")
            
            return state
        
        def review_node(state: Dict[str, Any]) -> Dict[str, Any]:
            """代码审查节点"""
            state['current_step'] = "代码审查"
            logger.info("🔍 开始代码审查...")
            
            try:
                result = self.agents['reviewer'].process(state['data'])
                result['success'] = True
                state['results']['review'] = result
                state['data'].update(result)
                state['completed_steps'].append('review')
                logger.info("✅ 代码审查完成")
            except Exception as e:
                state['errors'].append({
                    'step': 'review',
                    'error': str(e),
                    'timestamp': time.time()
                })
                state['results']['review'] = {'success': False, 'error': str(e)}
                logger.error(f"❌ 代码审查失败: {e}")
            
            return state
        
        def security_check_node(state: Dict[str, Any]) -> Dict[str, Any]:
            """安全检查节点"""
            state['current_step'] = "安全检查"
            logger.info("🔒 开始安全检查...")
            
            try:
                result = self.agents['security_checker'].process(state['data'])
                result['success'] = True
                state['results']['security_check'] = result
                state['data'].update(result)
                state['completed_steps'].append('security_check')
                logger.info("✅ 安全检查完成")
            except Exception as e:
                state['errors'].append({
                    'step': 'security_check',
                    'error': str(e),
                    'timestamp': time.time()
                })
                state['results']['security_check'] = {'success': False, 'error': str(e)}
                logger.error(f"❌ 安全检查失败: {e}")
            
            return state
        
        def optimization_node(state: Dict[str, Any]) -> Dict[str, Any]:
            """优化节点"""
            state['current_step'] = "性能优化"
            logger.info("⚡ 开始性能优化...")
            
            try:
                result = self.agents['optimizer'].process(state['data'])
                result['success'] = True
                state['results']['optimization'] = result
                state['data'].update(result)
                state['completed_steps'].append('optimization')
                logger.info("✅ 性能优化完成")
            except Exception as e:
                state['errors'].append({
                    'step': 'optimization',
                    'error': str(e),
                    'timestamp': time.time()
                })
                state['results']['optimization'] = {'success': False, 'error': str(e)}
                logger.error(f"❌ 性能优化失败: {e}")
            
            return state
        
        def documentation_node(state: Dict[str, Any]) -> Dict[str, Any]:
            """文档生成节点"""
            state['current_step'] = "文档生成"
            logger.info("📚 开始文档生成...")
            
            try:
                result = self.agents['documenter'].process(state['data'])
                result['success'] = True
                state['results']['documentation'] = result
                state['data'].update(result)
                state['completed_steps'].append('documentation')
                logger.info("✅ 文档生成完成")
            except Exception as e:
                state['errors'].append({
                    'step': 'documentation',
                    'error': str(e),
                    'timestamp': time.time()
                })
                state['results']['documentation'] = {'success': False, 'error': str(e)}
                logger.error(f"❌ 文档生成失败: {e}")
            
            return state
        
        def decide_next_after_testing(state: Dict[str, Any]) -> str:
            """决定测试后的下一步"""
            return "review"
        
        def decide_next_after_review(state: Dict[str, Any]) -> str:
            """决定审查后的下一步"""
            return "security_check"
        
        def decide_next_after_security(state: Dict[str, Any]) -> str:
            """决定安全检查后的下一步"""
            return "optimization"
        
        def decide_next_after_optimization(state: Dict[str, Any]) -> str:
            """决定性能优化后的下一步"""
            return "code_improvement"
        
        def code_improvement_node(state: Dict[str, Any]) -> Dict[str, Any]:
            """代码改进节点 - 根据反馈意见改进代码"""
            state['current_step'] = "代码改进"
            logger.info("🔧 开始代码改进...")
            
            try:
                # 收集所有反馈意见
                feedback = {
                    'testing_feedback': state['results'].get('testing', {}).get('feedback', []),
                    'review_feedback': state['results'].get('review', {}).get('feedback', []),
                    'security_feedback': state['results'].get('security_check', {}).get('feedback', []),
                    'optimization_feedback': state['results'].get('optimization', {}).get('feedback', [])
                }
                
                # 将反馈意见添加到数据中
                state['data']['improvement_feedback'] = feedback
                state['data']['original_code'] = state['results'].get('development', {}).get('code', '')
                
                # 调用开发工程师进行代码改进
                result = self.agents['developer'].process(state['data'])
                result['success'] = True
                result['improvement_type'] = 'feedback_based'
                state['results']['code_improvement'] = result
                state['data'].update(result)
                state['completed_steps'].append('code_improvement')
                logger.info("✅ 代码改进完成")
            except Exception as e:
                state['errors'].append({
                    'step': 'code_improvement',
                    'error': str(e),
                    'timestamp': time.time()
                })
                state['results']['code_improvement'] = {'success': False, 'error': str(e)}
                logger.error(f"❌ 代码改进失败: {e}")
            
            return state
        
        # 构建状态图
        workflow = StateGraph(WorkflowState)
        
        # 添加节点
        workflow.add_node("product_analysis", product_analysis_node)
        workflow.add_node("architecture_design", architecture_design_node)
        workflow.add_node("development", development_node)
        workflow.add_node("testing", testing_node)
        workflow.add_node("review", review_node)
        workflow.add_node("security_check", security_check_node)
        workflow.add_node("optimization", optimization_node)
        workflow.add_node("code_improvement", code_improvement_node)
        workflow.add_node("documentation", documentation_node)
        
        # 定义工作流 - 改为顺序执行避免并行汇聚问题
        workflow.set_entry_point("product_analysis")
        
        # 线性流程
        workflow.add_edge("product_analysis", "architecture_design")
        workflow.add_edge("architecture_design", "development")
        workflow.add_edge("development", "testing")
        
        # 顺序执行审查阶段 - 避免并行汇聚
        workflow.add_conditional_edges(
            "testing",
            decide_next_after_testing,
            {
                "review": "review"
            }
        )
        
        workflow.add_conditional_edges(
            "review",
            decide_next_after_review,
            {
                "security_check": "security_check"
            }
        )
        
        workflow.add_conditional_edges(
            "security_check",
            decide_next_after_security,
            {
                "optimization": "optimization"
            }
        )
        
        workflow.add_conditional_edges(
            "optimization",
            decide_next_after_optimization,
            {
                "code_improvement": "code_improvement"
            }
        )
        
        # 代码改进后生成文档
        workflow.add_edge("code_improvement", "documentation")
        
        # 结束
        workflow.add_edge("documentation", END)
        
        return workflow.compile()
    
    def execute(self, user_input: str, config: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
        """执行完整工作流"""
        logger.info("🚀 开始执行工作流...")
        start_time = time.time()
        
        # 初始化状态
        initial_state = {
            'user_input': user_input,
            'original_requirement': user_input,
            'data': {
                'user_input': user_input,
                'original_requirement': user_input
            },
            'results': {},
            'errors': [],
            'execution_log': [],
            'current_step': '',
            'completed_steps': []  # 确保是列表
        }
        
        try:
            # 执行状态图
            final_state = self.graph.invoke(initial_state)
            
            # 计算执行时间
            execution_time = time.time() - start_time
            
            # 获取实际执行的步骤数量
            total_steps_executed = len(final_state['completed_steps'])
            total_errors = len(final_state['errors'])
            
            # 生成执行摘要
            summary = {
                'total_steps': total_steps_executed,
                'successful_steps': total_steps_executed - total_errors,
                'failed_steps': total_errors,
                'execution_time': round(execution_time, 2),
                'success_rate': ((total_steps_executed - total_errors) / total_steps_executed * 100) if total_steps_executed > 0 else 0,
                'errors': total_errors
            }
            
            # 显示执行摘要
            logger.info(f"📊 执行摘要:")
            logger.info(f"   总步骤数: {summary['total_steps']}")
            logger.info(f"   成功步骤: {summary['successful_steps']}")
            logger.info(f"   失败步骤: {summary['failed_steps']}")
            logger.info(f"   执行时间: {summary['execution_time']}秒")
            logger.info(f"   成功率: {summary['success_rate']:.1f}%")
            
            return {
                'success': True,
                'results': final_state['results'],
                'summary': summary,
                'execution_log': final_state['execution_log'],
                'errors': final_state['errors']
            }
            
        except Exception as e:
            execution_time = time.time() - start_time
            logger.error(f"❌ 工作流执行失败: {e}")
            logger.error(f"⏱️ 执行时间: {round(execution_time, 2)}秒")
            return {
                'success': False,
                'error': str(e),
                'results': initial_state['results'],
                'execution_log': initial_state['execution_log'],
                'errors': initial_state['errors'],
                'execution_time': round(execution_time, 2)
            }
    
    def execute_partial(self, user_input: str, steps: List[str]) -> Dict[str, Any]:
        """执行部分工作流步骤"""
        logger.info(f"🔄 执行部分工作流: {', '.join(steps)}")
        start_time = time.time()
        
        # 初始化状态
        state = {
            'user_input': user_input,
            'original_requirement': user_input,
            'data': {},
            'results': {},
            'errors': [],
            'execution_log': [],
            'current_step': '',
            'completed_steps': []
        }
        
        results = {}
        
        # 按顺序执行指定步骤
        for step in steps:
            if step in self.agents:
                try:
                    logger.info(f"▶️ 执行步骤: {step}")
                    result = self.agents[step].process(state['data'])
                    state['results'][step] = result
                    state['data'].update(result)
                    state['completed_steps'].append(step)
                    results[step] = result
                except Exception as e:
                    logger.error(f"❌ 步骤 {step} 执行失败: {e}")
                    state['errors'].append({
                        'step': step,
                        'error': str(e),
                        'timestamp': time.time()
                    })
        
        # 计算执行时间
        execution_time = time.time() - start_time
        total_steps = len(steps)
        successful_steps = len(state['completed_steps'])
        failed_steps = len(state['errors'])
        
        # 显示执行摘要
        logger.info(f"📊 部分执行摘要:")
        logger.info(f"   计划步骤: {total_steps}")
        logger.info(f"   成功步骤: {successful_steps}")
        logger.info(f"   失败步骤: {failed_steps}")
        logger.info(f"   执行时间: {round(execution_time, 2)}秒")
        logger.info(f"   成功率: {(successful_steps / total_steps * 100):.1f}%" if total_steps > 0 else "0%")
        
        return {
            'success': len(state['errors']) == 0,
            'results': results,
            'errors': state['errors'],
            'execution_time': round(execution_time, 2),
            'summary': {
                'planned_steps': total_steps,
                'successful_steps': successful_steps,
                'failed_steps': failed_steps,
                'execution_time': round(execution_time, 2),
                'success_rate': (successful_steps / total_steps * 100) if total_steps > 0 else 0
            }
        }
    
    def get_agent_stats(self) -> Dict[str, Any]:
        """获取所有代理的统计信息"""
        stats = {}
        for name, agent in self.agents.items():
            stats[name] = agent.get_stats()
        return stats
    
    def reset_agents(self):
        """重置所有代理的状态"""
        for agent in self.agents.values():
            agent.reset_history()
        logger.info("🔄 已重置所有代理状态")

# 创建全局工作流实例
def create_workflow() -> CodeBuddyWorkflow:
    """创建CodeBuddy工作流实例"""
    return CodeBuddyWorkflow()

# 便捷函数
def quick_execute(user_input: str) -> Dict[str, Any]:
    """快速执行完整工作流"""
    workflow = create_workflow()
    return workflow.execute(user_input) 